#
#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
import json
import os
import sys
import time
import argparse
from collections import defaultdict

from api.db import LLMType
from api.db.services.llm_service import LLMBundle
from api.db.services.knowledgebase_service import KnowledgebaseService
from api import settings
from api.utils import get_uuid
from rag.nlp import tokenize, search
from ranx import evaluate
from ranx import Qrels, Run
import pandas as pd
from tqdm import tqdm

global max_docs
max_docs = sys.maxsize


# 定义Benchmark类，用于对指定知识库进行基准测试
class Benchmark:
    # 初始化方法，接收知识库ID作为参数
    def __init__(self, kb_id):
        # 获取知识库信息
        self.kb_id = kb_id
        e, self.kb = KnowledgebaseService.get_by_id(kb_id)
        # 设置相似度阈值
        self.similarity_threshold = self.kb.similarity_threshold
        # 设置向量相似度权重
        self.vector_similarity_weight = self.kb.vector_similarity_weight
        # 初始化嵌入模型
        self.embd_mdl = LLMBundle(self.kb.tenant_id, LLMType.EMBEDDING, llm_name=self.kb.embd_id, lang=self.kb.language)
        # 初始化租户ID为空字符串
        self.tenant_id = ''
        # 初始化索引名称为空字符串
        self.index_name = ''
        # 标记索引是否已初始化
        self.initialized_index = False

    # 私有方法_get_retrieval，用于执行检索并返回结果
    def _get_retrieval(self, qrels):
        # Need to wait for the ES and Infinity index to be ready
        time.sleep(20)
        # 使用defaultdict存储运行结果
        run = defaultdict(dict)
        # 提取查询列表
        query_list = list(qrels.keys())
        # 遍历每个查询
        for query in query_list:
            # 执行检索操作
            ranks = settings.retrievaler.retrieval(query, self.embd_mdl, self.tenant_id, [self.kb.id], 1, 30,
                                            0.0, self.vector_similarity_weight)
            # 如果没有检索到结果，则从qrels中删除该查询
            if len(ranks["chunks"]) == 0:
                print(f"deleted query: {query}")
                del qrels[query]
                continue
            # 处理每个检索到的chunk，去除vector字段并将相似度存入run字典
            for c in ranks["chunks"]:
                c.pop("vector", None)
                # 将检索结果整理为标准格式，便于后续评估
                run[query][c["chunk_id"]] = c["similarity"]
        # 返回检索结果
        return run

    # 为给定的文档列表生成内容的嵌入向量，并将结果附加到文档字典中
    def embedding(self, docs):
        # 提取所有文档的内容部分用于编码
        texts = [d["content_with_weight"] for d in docs]
        # 使用嵌入模型对文本进行编码，得到嵌入向量和异常信息
        embeddings, _ = self.embd_mdl.encode(texts)
        # 确保返回的嵌入数量与文档数量一致
        assert len(docs) == len(embeddings)
        # 初始化向量维度大小
        vector_size = 0
        # 遍历每个文档及其对应的嵌入向量
        for i, d in enumerate(docs):
            v = embeddings[i]
            # 记录当前向量的维度
            vector_size = len(v)
            # 将向量添加到文档中，键为"q_{维度}_vec"
            d["q_%d_vec" % len(v)] = v
        # 返回更新后的文档列表和向量维度
        return docs, vector_size

    # 初始化索引的方法，确保目标索引存在并清空旧数据（如果已初始化）
    def init_index(self, vector_size: int):
        # 如果已经初始化过索引，则直接返回
        if self.initialized_index:
            return
        # 检查是否存在同名索引，如果存在则删除
        if settings.docStoreConn.indexExist(self.index_name, self.kb_id):
            settings.docStoreConn.deleteIdx(self.index_name, self.kb_id)
        # 创建新的索引，指定知识库ID和向量维度
        settings.docStoreConn.createIdx(self.index_name, self.kb_id, vector_size)
        # 标记索引已初始化
        self.initialized_index = True

    # 构建MS MARCO风格的索引，并返回qrels（相关性标注）和texts（文档内容）
    def ms_marco_index(self, file_path, index_name):
        # 初始化qrels为defaultdict(dict)，用于存储查询与文档的相关性关系
        qrels = defaultdict(dict)
        # 初始化texts为defaultdict(dict)，用于存储文档ID到文本内容的映射
        texts = defaultdict(dict)
        # 初始化已处理文档数量计数器
        docs_count = 0
        # 初始化当前批次待处理的文档列表
        docs = []
        # 获取文件路径下的所有文件，并按名称排序
        filelist = sorted(os.listdir(file_path))

        # 遍历每个文件
        for fn in filelist:
            # 如果已达到最大文档数限制，则跳出循环
            if docs_count >= max_docs:
                break
            # 跳过非parquet格式的文件
            if not fn.endswith(".parquet"):
                continue
            # 读取parquet文件为DataFrame格式
            data = pd.read_parquet(os.path.join(file_path, fn))
            # 遍历DataFrame中的每一行
            for i in tqdm(range(len(data)), colour="green", desc="Tokenizing:" + fn):
                # 如果已达到最大文档数限制，则跳出内层循环
                if docs_count >= max_docs:
                    break
                # 提取当前行的查询语句
                query = data.iloc[i]['query']
                # 遍历passages中的每一个段落及其相关性标签
                for rel, text in zip(data.iloc[i]['passages']['is_selected'], data.iloc[i]['passages']['passage_text']):
                    # 创建一个文档字典结构
                    d = {
                        "id": get_uuid(),
                        "kb_id": self.kb.id,
                        "docnm_kwd": "xxxxx",
                        "doc_id": "ksksks"
                    }
                    # 对文本进行分词处理并更新文档d
                    tokenize(d, text, "english")
                    # 将文档添加到docs列表中
                    docs.append(d)
                    # 存储文档ID到文本的映射
                    texts[d["id"]] = text
                    # 记录查询与文档的相关性评分
                    qrels[query][d["id"]] = int(rel)
                # 如果当前批次的文档数量达到32，则执行批量处理
                if len(docs) >= 32:
                    # 对当前批次文档进行嵌入编码
                    docs_count += len(docs)
                    docs, vector_size = self.embedding(docs)
                    # 初始化或重建向量索引
                    self.init_index(vector_size)
                    # 插入文档到索引中
                    settings.docStoreConn.insert(docs, self.index_name, self.kb_id)
                    # 清空当前批次文档列表
                    docs = []

        # 处理剩余未插入的文档
        if docs:
            # 对剩余文档进行嵌入编码
            docs, vector_size = self.embedding(docs)
            # 初始化或重建向量索引
            self.init_index(vector_size)
            # 插入剩余文档到索引中
            settings.docStoreConn.insert(docs, self.index_name, self.kb_id)
        # 返回qrels（查询-文档相关性）和texts（文档内容）
        return qrels, texts

    # 构建TriviaQA格式的索引，并返回相关性标注(qrels)和文档内容(texts)
    def trivia_qa_index(self, file_path, index_name):
        # 初始化qrels字典，存储查询与文档的相关性关系
        qrels = defaultdict(dict)
        # 初始化texts字典，存储文档ID到文本内容的映射
        texts = defaultdict(dict)
        # 初始化已处理文档数量计数器
        docs_count = 0
        # 初始化当前批次待处理的文档列表
        docs = []
        # 获取文件路径下的所有文件，并按名称排序
        filelist = sorted(os.listdir(file_path))
        # 遍历每个文件
        for fn in filelist:
            # 如果已达到最大文档数限制，则跳出循环
            if docs_count >= max_docs:
                break
            # 跳过非parquet格式的文件
            if not fn.endswith(".parquet"):
                continue
            # 读取parquet文件为DataFrame格式
            data = pd.read_parquet(os.path.join(file_path, fn))
            # 遍历DataFrame中的每一行
            for i in tqdm(range(len(data)), colour="green", desc="Indexing:" + fn):
                # 如果已达到最大文档数限制，则跳出内层循环
                if docs_count >= max_docs:
                    break
                # 提取当前行的问题作为查询语句
                query = data.iloc[i]['question']
                # 遍历搜索结果及其排名
                for rel, text in zip(data.iloc[i]["search_results"]['rank'],
                                     data.iloc[i]["search_results"]['search_context']):
                    # 创建一个文档字典结构
                    d = {
                        "id": get_uuid(),
                        "kb_id": self.kb.id,
                        "docnm_kwd": "xxxxx",
                        "doc_id": "ksksks"
                    }
                    # 对文本进行分词处理并更新文档d
                    tokenize(d, text, "english")
                    # 将文档添加到docs列表中
                    docs.append(d)
                    # 存储文档ID到文本的映射
                    texts[d["id"]] = text
                    # 记录查询与文档的相关性评分
                    qrels[query][d["id"]] = int(rel)
                # 如果当前批次的文档数量达到32，则执行批量处理
                if len(docs) >= 32:
                    # 更新总文档计数
                    docs_count += len(docs)
                    # 对当前批次文档进行嵌入编码
                    docs, vector_size = self.embedding(docs)
                    # 初始化或重建向量索引
                    self.init_index(vector_size)
                    # 插入文档到索引中
                    settings.docStoreConn.insert(docs,self.index_name)
                    # 清空当前批次文档列表
                    docs = []

        # 对剩余未插入的文档进行嵌入编码
        docs, vector_size = self.embedding(docs)
        # 初始化或重建向量索引
        self.init_index(vector_size)
        # 插入剩余文档到索引中
        settings.docStoreConn.insert(docs, self.index_name)
        # 返回qrels（查询-文档相关性）和texts（文档内容）
        return qrels, texts

    # 构建MIRACL格式的索引，并返回相关性标注(qrels)和文档内容(texts)
    def miracl_index(self, file_path, corpus_path, index_name):
        # 初始化corpus_total字典，存储文档ID到文本的映射
        corpus_total = {}
        # 遍历corpus目录下的所有文件
        for corpus_file in os.listdir(corpus_path):
            # 读取JSON格式的语料文件
            tmp_data = pd.read_json(os.path.join(corpus_path, corpus_file), lines=True)
            # 遍历每条语料记录
            for index, i in tmp_data.iterrows():
                # 建立docid到text的映射
                corpus_total[i['docid']] = i['text']

        # 初始化topics_total字典，存储查询ID到问题的映射
        topics_total = {}
        # 遍历topics目录下的所有文件
        for topics_file in os.listdir(os.path.join(file_path, 'topics')):
            # 跳过测试用的topic文件
            if 'test' in topics_file:
                continue
            # 读取TSV格式的主题文件
            tmp_data = pd.read_csv(os.path.join(file_path, 'topics', topics_file), sep='\t', names=['qid', 'query'])
            # 遍历每条主题记录
            for index, i in tmp_data.iterrows():
                # 建立qid到query的映射
                topics_total[i['qid']] = i['query']

        # 初始化qrels字典，存储查询与文档的相关性关系
        qrels = defaultdict(dict)
        # 初始化texts字典，存储文档ID到文本内容的映射
        texts = defaultdict(dict)
        # 初始化已处理文档数量计数器
        docs_count = 0
        # 初始化当前批次待处理的文档列表
        docs = []
        # 遍历qrels目录下的所有文件
        for qrels_file in os.listdir(os.path.join(file_path, 'qrels')):
            # 跳过测试用的qrel文件
            if 'test' in qrels_file:
                continue
            # 如果已达到最大文档数限制，则跳出循环
            if docs_count >= max_docs:
                break

            # 读取相关性文件
            tmp_data = pd.read_csv(os.path.join(file_path, 'qrels', qrels_file), sep='\t',
                                   names=['qid', 'Q0', 'docid', 'relevance'])
            # 遍历每条相关性记录
            for i in tqdm(range(len(tmp_data)), colour="green", desc="Indexing:" + qrels_file):
                # 如果已达到最大文档数限制，则跳出内层循环
                if docs_count >= max_docs:
                    break
                # 根据qid获取对应的查询问题
                query = topics_total[tmp_data.iloc[i]['qid']]
                # 根据docid获取对应的文档内容
                text = corpus_total[tmp_data.iloc[i]['docid']]
                # 获取相关性分数
                rel = tmp_data.iloc[i]['relevance']
                # 创建一个文档字典结构
                d = {
                    "id": get_uuid(),
                    "kb_id": self.kb.id,
                    "docnm_kwd": "xxxxx",
                    "doc_id": "ksksks"
                }
                # 对文本进行分词处理并更新文档d
                tokenize(d, text, 'english')
                # 将文档添加到docs列表中
                docs.append(d)
                # 存储文档ID到文本的映射
                texts[d["id"]] = text
                # 记录查询与文档的相关性评分
                qrels[query][d["id"]] = int(rel)
                # 如果当前批次的文档数量达到32，则执行批量处理
                if len(docs) >= 32:
                    # 更新总文档计数
                    docs_count += len(docs)
                    # 对当前批次文档进行嵌入编码
                    docs, vector_size = self.embedding(docs)
                    # 初始化或重建向量索引
                    self.init_index(vector_size)
                    # 插入文档到索引中
                    settings.docStoreConn.insert(docs, self.index_name)
                    # 清空当前批次文档列表
                    docs = []

        # 对剩余未插入的文档进行嵌入编码
        docs, vector_size = self.embedding(docs)
        # 初始化或重建向量索引
        self.init_index(vector_size)
        # 插入剩余文档到索引中
        settings.docStoreConn.insert(docs, self.index_name)
        # 返回qrels（查询-文档相关性）和texts（文档内容）
        return qrels, texts

    # 保存测试结果与相关性评分，并输出每个查询的ndcg@10得分及对应的文档内容
    def save_results(self, qrels, run, texts, dataset, file_path):
        # 初始化一个列表用于存储每个查询的结果及其评估得分
        keep_result = []
        # 获取所有查询的键值列表
        run_keys = list(run.keys())
        # 遍历每个查询并计算其ndcg@10得分
        for run_i in tqdm(range(len(run_keys)), desc="Calculating ndcg@10 for single query"):
            key = run_keys[run_i]
            # 对每个查询计算评估结果并保存
            keep_result.append({'query': key, 'qrel': qrels[key], 'run': run[key],
                                'ndcg@10': evaluate({key: qrels[key]}, {key: run[key]}, "ndcg@10")})
        # 按照ndcg@10得分对查询结果进行排序（从低到高）
        keep_result = sorted(keep_result, key=lambda kk: kk['ndcg@10'])
        # 打开或创建md文件用于写入评估结果报告
        with open(os.path.join(file_path, dataset + 'result.md'), 'w', encoding='utf-8') as f:
            # 写入标题：每查询得分
            f.write('## Score For Every Query\n')
            # 遍历每个查询的评估结果
            for keep_result_i in keep_result:
                # 写入当前查询及其ndcg@10得分
                f.write('### query: ' + keep_result_i['query'] + ' ndcg@10:' + str(keep_result_i['ndcg@10']) + '\n')
                # 提取当前查询的运行结果并按相似度得分排序
                scores = [[i[0], i[1]] for i in keep_result_i['run'].items()]
                scores = sorted(scores, key=lambda kk: kk[1])
                # 取出得分最高的前10个文档片段
                for score in scores[:10]:
                    # 写入文档内容和对应的相关性标签
                    f.write('- text: ' + str(texts[score[0]]) + '\t qrel: ' + str(score[1]) + '\n')
        # 将qrels数据以JSON格式保存到指定路径
        json.dump(qrels, open(os.path.join(file_path, dataset + '.qrels.json'), "w+", encoding='utf-8'), indent=2)
        # 将run数据（检索结果）以JSON格式保存到指定路径
        json.dump(run, open(os.path.join(file_path, dataset + '.run.json'), "w+", encoding='utf-8'), indent=2)
        # 输出保存提示信息
        print(os.path.join(file_path, dataset + '_result.md'), 'Saved!')

    # 该方法使实例可像函数一样被调用，根据不同的数据集名称执行对应的任务流程
    def __call__(self, dataset, file_path, miracl_corpus=''):
        # 如果数据集是ms_marco_v1.1，则执行MS MARCO格式的处理流程
        if dataset == "ms_marco_v1.1":
            # 设置租户ID为benchmark_ms_marco_v11
            self.tenant_id = "benchmark_ms_marco_v11"
            # 构建索引名称
            self.index_name = search.index_name(self.tenant_id)
            # 构建索引并获取qrels和texts
            qrels, texts = self.ms_marco_index(file_path, "benchmark_ms_marco_v1.1")
            # 执行检索任务
            run = self._get_retrieval(qrels)
            # 打印评估结果（ndcg@10, map@5, mrr@10）
            print(dataset, evaluate(Qrels(qrels), Run(run), ["ndcg@10", "map@5", "mrr@10"]))
            # 保存评估结果到文件
            self.save_results(qrels, run, texts, dataset, file_path)
        # 如果数据集是trivia_qa，则执行TriviaQA格式的处理流程
        if dataset == "trivia_qa":
            # 设置租户ID为benchmark_trivia_qa
            self.tenant_id = "benchmark_trivia_qa"
            # 构建索引名称
            self.index_name = search.index_name(self.tenant_id)
            # 构建索引并获取qrels和texts
            qrels, texts = self.trivia_qa_index(file_path, "benchmark_trivia_qa")
            # 执行检索任务
            run = self._get_retrieval(qrels)
            # 打印评估结果（ndcg@10, map@5, mrr@10）
            print(dataset, evaluate(Qrels(qrels), Run(run), ["ndcg@10", "map@5", "mrr@10"]))
            # 保存评估结果到文件
            self.save_results(qrels, run, texts, dataset, file_path)
        # 如果数据集是miracl，则处理多语言MIRACL数据集
        if dataset == "miracl":
            # 遍历所有支持的语言版本
            for lang in ['ar', 'bn', 'de', 'en', 'es', 'fa', 'fi', 'fr', 'hi', 'id', 'ja', 'ko', 'ru', 'sw', 'te', 'th',
                         'yo', 'zh']:
                # 检查主数据目录是否存在
                if not os.path.isdir(os.path.join(file_path, 'miracl-v1.0-' + lang)):
                    print('Directory: ' + os.path.join(file_path, 'miracl-v1.0-' + lang) + ' not found!')
                    continue
                # 检查qrels目录是否存在
                if not os.path.isdir(os.path.join(file_path, 'miracl-v1.0-' + lang, 'qrels')):
                    print('Directory: ' + os.path.join(file_path, 'miracl-v1.0-' + lang, 'qrels') + 'not found!')
                    continue
                # 检查topics目录是否存在
                if not os.path.isdir(os.path.join(file_path, 'miracl-v1.0-' + lang, 'topics')):
                    print('Directory: ' + os.path.join(file_path, 'miracl-v1.0-' + lang, 'topics') + 'not found!')
                    continue
                # 检查语料库目录是否存在
                if not os.path.isdir(os.path.join(miracl_corpus, 'miracl-corpus-v1.0-' + lang)):
                    print('Directory: ' + os.path.join(miracl_corpus, 'miracl-corpus-v1.0-' + lang) + ' not found!')
                    continue
                # 设置租户ID包含语言标识
                self.tenant_id = "benchmark_miracl_" + lang
                # 构建索引名称
                self.index_name = search.index_name(self.tenant_id)
                # 重置索引初始化标志位
                self.initialized_index = False
                # 构建当前语言的索引并获取qrels和texts
                qrels, texts = self.miracl_index(os.path.join(file_path, 'miracl-v1.0-' + lang),
                                                 os.path.join(miracl_corpus, 'miracl-corpus-v1.0-' + lang),
                                                 "benchmark_miracl_" + lang)
                # 执行检索任务
                run = self._get_retrieval(qrels)
                # 打印评估结果（ndcg@10, map@5, mrr@10）
                print(dataset, evaluate(Qrels(qrels), Run(run), ["ndcg@10", "map@5", "mrr@10"]))
                # 保存评估结果到文件
                self.save_results(qrels, run, texts, dataset, file_path)


# 当前模块作为主程序运行时才执行以下内容
if __name__ == '__main__':
    # 打印欢迎信息
    print('*****************RAGFlow Benchmark*****************')
    parser = argparse.ArgumentParser(usage="benchmark.py <max_docs> <kb_id> <dataset> <dataset_path> [<miracl_corpus_path>])", description='RAGFlow Benchmark')
    parser.add_argument('max_docs', metavar='max_docs', type=int, help='max docs to evaluate')
    # 添加 kb_id 参数：知识库ID（字符串）
    parser.add_argument('kb_id', metavar='kb_id', help='knowledgebase id')
    # 添加 dataset参数
    parser.add_argument('dataset', metavar='dataset', help='dataset name, shall be one of ms_marco_v1.1(https://huggingface.co/datasets/microsoft/ms_marco), trivia_qa(https://huggingface.co/datasets/mandarjoshi/trivia_qa>), miracl(https://huggingface.co/datasets/miracl/miracl')
    parser.add_argument('dataset_path', metavar='dataset_path', help='dataset path')
    # 添加可选的 miracl_corpus_path 参数：仅当使用 miracl 数据集时需要
    parser.add_argument('miracl_corpus_path', metavar='miracl_corpus_path', nargs='?', default="", help='miracl corpus path. Only needed when dataset is miracl')

    # 解析命令行输入为 args 对象
    args = parser.parse_args()
    # 将命令行参数赋值给变量
    max_docs = args.max_docs
    kb_id = args.kb_id
    # 初始化 Benchmark 实例，传入 kb_id
    ex = Benchmark(kb_id)

    # 获取数据集名称和路径
    dataset = args.dataset
    dataset_path = args.dataset_path

    # 根据数据集类型调用 Benchmark.__call__()
    if dataset == "ms_marco_v1.1" or dataset == "trivia_qa":
        # 调用 Benchmark 实例，只传入两个参数
        ex(dataset, dataset_path)
    elif dataset == "miracl":
        # 检查是否提供了 miracl_corpus_path 参数
        if len(args) < 5:
            print('Please input the correct parameters!')
            exit(1)
        # 获取 miracl_corpus_path
        miracl_corpus_path = args[4]
        # 调用 Benchmark 实例，并传入三个参数
        ex(dataset, dataset_path, miracl_corpus=args.miracl_corpus_path)
    else:
        # 如果数据集不支持，则输出错误信息
        print("Dataset: ", dataset, "not supported!")
